#
# 4.3 自定义层
# 通过对 tf.keras.layers.Layer 进行子类化并实现以下方法来创建自定义层：
#
# _init_: (可选)定义该层要使用的子层
#
# build：创建层的权重。使用 add_weight 方法添加权重。
#
# call：定义前向传播。
#
# compute_output_shape：指定在给定输入形状的情况下如何计算层的输出形状。
#
# 可选，可以通过实现 get_config 方法和 from_config 类方法序列化层。

import tensorflow as tf
from tensorflow import keras

from tensorflow.keras import layers

import  numpy as np

print(tf.__version__)  # 2.1.0
print(keras.__version__)  # 2.2.4-tf


# 自定义层
class MyLayer(layers.Layer):
    def __init__(self, output_dim, **kwargs):
        self.output_dim = output_dim
        super(MyLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        shape = tf.TensorShape((input_shape[1], self.output_dim))
        self.kernel = self.add_weight(name="kernel1", shape=shape,
                                      initializer='uniform', trainable=True)

        super(MyLayer, self).build(input_shape)


    def call(self, inputs):

        return tf.matmul(inputs,self.kernel)


    def compute_output_shape(self, input_shape):
        shape= tf.TensorShape(input_shape).as_list()
        shape[-1] = self.output_dim

        return  tf.TensorShape(shape)


    def get_config(self):

        base_config = super(MyLayer,self).get_config()

        base_config['output_dim'] = self.output_dim

        return  base_config

    @classmethod
    def from_config(cls, config):
        return  cls(**config)



# use user define layers create model

model = tf.keras.Sequential([
    MyLayer(10),
    layers.Activation('softmax')
])


model.compile(optimizer= tf.keras.optimizers.RMSprop(0.001),
              loss=tf.keras.losses.categorical_crossentropy,
              metrics=['accuracy']
        )

# 模型评估
test_x = np.random.random((1000, 72))
test_y = np.random.random((1000, 10))
test_data = tf.data.Dataset.from_tensor_slices((test_x, test_y))
test_data = test_data.batch(32).repeat()


model.fit(test_x,test_y,epochs=5,batch_size=16)
